# -*- coding:utf-8 -*-
from sqlalchemy import create_engine
import pandas as pd
import tushare as ts
import datetime, time


class Stock:
    def __init__(self):
        self.engine = create_engine('mysql+mysqldb://root:root@127.0.0.1/stock?charset=utf8')

    # 获取沪深上市公司的基本情况
    def get_stock_list(self):
        sql = 'SELECT table_name FROM information_schema.TABLES WHERE table_name ="stocks"'
        df = pd.read_sql_query(sql, self.engine)
        if df.empty:
            print 'stocks表不存在,请先运行createtable.py脚本'
            exit()
        sql = 'truncate table stocks '
        self.engine.execute(sql)  # 清空数据
        df = ts.get_stock_basics()
        df.to_sql('stocks', self.engine, if_exists='append')
        print '沪深上市公司抓取完成'
        exit()

    def get_stock_date(self):
        sql = 'select MAX(date) as date,code from stock_data   GROUP BY code'
        stocks = self.engine.execute(sql)
        dic = {}
        for item in stocks:
            dic[item.code] = item.date.strftime('%Y-%m-%d')
        return dic

    def get_stock_right_date(self):
        sql = 'select MAX(date) as date,code from stock_right_data   GROUP BY code'
        stocks = self.engine.execute(sql)
        dic = {}
        for item in stocks:
            dic[item.code] = item.date.strftime('%Y-%m-%d')
        return dic

    # 获取个股历史交易记录---除权记录
    def get_stock_hist_data(self, code, start, end):
        if code == "":
            print '代码不能为空'
            exit()
        if start < end:
            df = ts.get_hist_data(code, start=start, end=end)
            if df is None:
                pass
            else:
                df['code'] = pd.Series([])
                df = df.fillna(code)
                df.to_sql('stock_data', self.engine, if_exists='append')
                print '代码' + code + '抓取完成(除权)'
        else:
            print code + '已经抓取过了!!!'

    # 获取个股历史交易记录----前复权记录
    def get_stock_h_data(self, code, start, end):
        if code == "":
            print '代码不能为空'
            exit()
        if start < end:
            df = ts.get_h_data(code, start=start, end=end)
            if df is None:
                pass
            else:
                df['code'] = pd.Series([])
                df = df.fillna(code)
                df.to_sql('stock_right_data', self.engine, if_exists='append')
                print '代码' + code + '抓取完成(前复权)'
        else:
            print code + '已经抓取过了!!!'

    def get_stock_from_table(self):
        sql = 'select code from stocks where id<1000 order by id'
        stocks = self.engine.execute(sql)
        return stocks.fetchall()

    def get_all_stock_hist(self, right=None):
        stocks = self.get_stock_from_table()
        if right is None:
            codes = self.get_stock_date()
        else:
            codes = self.get_stock_right_date()
        end = datetime.date.today().strftime('%Y-%m-%d')  # 获得今天的日期
        for code in stocks:
            if codes.has_key(code.code):
                start = codes[code[0].encode('utf-8')]
                start = datetime.datetime.strptime(start, '%Y-%m-%d')
                start = start + datetime.timedelta(days=1)
                start = start.strftime('%Y-%m-%d')
            else:
                start = '2017-01-01'
            if right is None:
                self.get_stock_hist_data(code[0].encode('utf-8'), start, end)
            else:
                self.get_stock_h_data(code[0].encode('utf-8'), start, end)
        print '所有数据都抓完了!!!!'

    # 获取当天所有股票的数据，节假日，即为上一交易日
    def get_toady_stock_all(self):

        if time.localtime(time.time()).tm_hour >= 16:
            date = datetime.date.today().strftime('%Y-%m-%d')
            sql = "select date,code from stock_right_data where date = '" + date + "'"
            result = self.engine.execute(sql)
            rows = result.fetchall()
            if rows:
                print date + '期数据已经存在'
            else:
                df = ts.get_today_all()
                df.rename(columns={'trade': 'close'}, inplace=True)
                df.set_index('code', inplace=True)
                del df['name']
                del df['changepercent']
                del df['settlement']
                del df['turnoverratio']
                del df['per']
                del df['pb']
                del df['mktcap']
                del df['nmc']
                df['date'] = pd.Series([])
                df = df.fillna(datetime.date.today().strftime('%Y-%m-%d'))
                df.to_sql('stock_right_data', self.engine, if_exists='append')
                print date + '期数据写入成功',
        else:
            print '还没有收盘'


stock = Stock()
stock.get_all_stock_hist('ok')  # 加参数为前复权

# stock.get_stock_list() # 获取沪深上市公司的基本情况

# stock.get_toady_stock_all() # 今天所有股票数据
